import os
import random

from PIL import Image, ImageDraw
from tqdm import tqdm


# 生成随机颜色
def random_color():
    return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))


# 绘制随机直线
def draw_random_lines(draw, width, height, num_lines):
    for _ in range(num_lines):
        start_point = (random.randint(0, width), random.randint(0, height))
        end_point = (random.randint(0, width), random.randint(0, height))
        draw.line([start_point, end_point], fill=random_color(), width=random.randint(5, 10))


# 绘制随机曲线（贝塞尔曲线）
# 绘制二次贝塞尔曲线
def bezier_curve(p0, p1, p2, t):
    """计算贝塞尔曲线上的点，t 介于 [0, 1]"""
    x = (1 - t) ** 2 * p0[0] + 2 * (1 - t) * t * p1[0] + t ** 2 * p2[0]
    y = (1 - t) ** 2 * p0[1] + 2 * (1 - t) * t * p1[1] + t ** 2 * p2[1]
    return (int(x), int(y))


# 绘制随机曲线
def draw_random_curves(draw, width, height, num_curves):
    for _ in range(num_curves):
        # 随机生成控制点
        p0 = (random.randint(0, width), random.randint(0, height))  # 起点
        p1 = (random.randint(0, width), random.randint(0, height))  # 控制点
        p2 = (random.randint(0, width), random.randint(0, height))  # 终点

        # 生成曲线上的点
        points = [bezier_curve(p0, p1, p2, t / 20) for t in range(21)]

        # 绘制曲线
        draw.line(points, fill=random_color(), width=random.randint(1, 5))


# 主函数
def create_image_with_lines(width_range, height_range, num_lines, num_curves):
    width = random.randint(*width_range)
    height = random.randint(*height_range)

    background_color = random_color()
    image = Image.new("RGB", (width, height), background_color)
    draw = ImageDraw.Draw(image)

    # 绘制随机直线和曲线
    draw_random_lines(draw, width, height, num_lines)
    draw_random_curves(draw, width, height, num_curves)
    return image


if __name__ == '__main__':
    out_dir = ''
    for i in tqdm(range(10000)):
        img = create_image_with_lines((512, 1024), (512, 1024), 20, 20)
        img.save(os.path.join(out_dir, f'{i}.png'))
